<template>
    <el-container style="padding: 20px;">
      <el-main>
        <h1>电磁环境感知模型训练与可视化平台</h1>
  
        <!-- 模型训练 -->
        <h3 style="margin-bottom: 10px;">模型训练</h3>
  
        <!-- 左右卡片布局 -->
        <el-row :gutter="20" style="margin-bottom: 30px;">
          <!-- 左侧卡片：上传数据集 & 参数配置 -->
          <el-col :span="12">
            <el-card style="height: 100%;">
              <div slot="header">
                <strong>上传数据集与参数配置</strong>
              </div>
              <div style="margin-bottom: 20px;">
                <strong>导入训练数据集</strong>
                <el-upload
                  drag
                  action="http://127.0.0.1:5000/api/dcgz/upload"
                  :before-upload="beforeUpload"
                  :on-success="handleSuccess"
                  multiple
                  :data="{ type: uploadType,username:username }"
                >
                  <i class="el-icon-upload"></i>
                  <div class="el-upload__text">将数据集压缩包拖到此处，或<em>点击上传</em></div>
                </el-upload>
                <p v-if="datasetPath" style="color: green; margin-top: 10px;">
                  ✅ 已上传路径：{{ datasetPath }}
                </p>
              </div>
  
              <!-- 参数配置 -->
              <el-form label-width="150px">
                <el-form-item label="模型类型">
                  YOLOV5
                </el-form-item>
  
                <el-form-item label="训练轮数 (Epochs)">
                  <el-input-number v-model="config.epochs" :min="1" :max="100"></el-input-number>
                </el-form-item>
  
                <el-form-item label="批次大小 (Batch Size)">
                  <el-input-number v-model="config.batchSize" :min="1" :max="64"></el-input-number>
                </el-form-item>
  
                <el-button type="primary" @click="startTraining">开始训练</el-button>
              </el-form>
            </el-card>
          </el-col>
  
          <!-- 右侧卡片：训练状态/结果 -->
          <el-col :span="12">
  <el-card style="height: 100%;">
    <div slot="header">
      <strong>训练状态</strong>
    </div>

    <div v-if="trainingStatus === 'idle'" style="text-align: center; padding: 60px 20px; color: #999;">
      <p>等待开始训练...</p>
    </div>

    <div v-else-if="trainingStatus === 'running'" style="padding: 20px;">
      <h4>训练进行中...</h4>
      
    </div>

    <div v-else-if="trainingStatus === 'completed'" style="padding: 15px;">
      <h4>训练完成</h4>
      <!-- 展示训练结果图片 -->
      <div v-if="trainingResult" style="margin-top: 10px;">

<!-- 横向排列三个图像 -->
<div v-if="trainingResult && trainingResult.results_png" class="image-container">
      <!-- 循环渲染每张图片 -->
      <div v-for="(img, index) in images" :key="index" class="image-item">
        <el-card shadow="hover" style="margin-bottom: 20px;">
          <!-- 使用 el-image 实现点击放大 -->
          <el-image
            :src="'data:image/png;base64,' + img.image_base64"
            :preview-src-list="allBase64Images"
            class="image-preview"
            fit="cover"
            lazy
            scroll-container=".image-container"
          >
          </el-image>
          <div style="text-align: center; padding-top: 10px;">
            <span>{{ getTitle(index) }}</span>
          </div>
        </el-card>
      </div>
    </div>
      </div>
    </div>

    <div v-else-if="trainingStatus === 'error'" style="padding: 20px; color: red;">
      <p>训练出错，请查看日志或重启服务后重试。</p>
    </div>
  </el-card>
 
</el-col>
</el-row><!-- 图片预览弹窗 -->

  
        <!-- 模型预测 -->
        <h3 style="margin-bottom: 10px;">模型预测</h3>
  
        <!-- 左右卡片布局 -->
        <el-row :gutter="20" style="margin-bottom: 30px;">
          <!-- 左侧卡片：上传数据集 -->
          <el-col :span="12">
            <el-card style="height: 100%;">
              <div slot="header">
                <strong>上传预测数据集</strong>
              </div>
              <div style="margin-bottom: 20px;">
                <el-upload
                  drag
                  action="http://127.0.0.1:5000/api/dcgz/upload"
                  :before-upload="beforeUpload"
                  :on-success="handlePredictSuccess"
                  multiple
                  :data="{ type: predictUploadType,username:username}"
                >
                  <i class="el-icon-upload"></i>
                  <div class="el-upload__text">将预测数据集压缩包拖到此处，或<em>点击上传</em></div>
                </el-upload>
                <div style="margin-top: 20px;">
        <label style="font-weight: bold; display: block; margin-bottom: 8px;">选择模型：</label>
        <el-select v-model="selectedModel" placeholder="请选择模型" style="width: 100%;">
          <el-option
            v-for="model in availableModels"
            :key="model"
            :label="model"
            :value="model"
          ></el-option>
        </el-select>
      </div>

                <p v-if="predictDatasetPath" style="color: green; margin-top: 10px;">
                  ✅ 已上传路径：{{ predictDatasetPath }}
                </p>
              </div>
              <el-button type="primary" @click="startPredict">开始预测</el-button>
            </el-card>
          </el-col>
  
          <!-- 右侧卡片：预测状态/结果 -->
          <el-col :span="12">
            <el-card style="height: 100%;">
              <div slot="header">
                <strong>预测状态</strong>
              </div>
  
              <div v-if="predictStatus === 'idle'" style="text-align: center; padding: 60px 20px; color: #999;">
                <p>等待开始预测...</p>
              </div>
  
              <div v-else-if="predictStatus === 'running'" style="padding: 20px;">
                <h4>预测进行中...</h4>
                <el-progress :percentage="predictProgress" status="active"></el-progress>
              </div>
  
              <div v-else-if="predictStatus === 'completed'" style="padding: 20px;">
                <h4>预测完成</h4>
                <p>识别准确率：<strong style="color: #409EFF">{{ accuracy }}%</strong></p>
              </div>
  
              <div v-else-if="predictStatus === 'error'" style="padding: 20px; color: red;">
                <p>预测出错，请查看日志或重启服务后重试。</p>
              </div>
            </el-card>
          </el-col>
        </el-row>
  
        <!-- 结果可视化卡片 -->
      </el-main>
    </el-container>
  </template>
  
  <script>
  import axios from 'axios';
  
  export default {
    data() {
      return {
       
        username: '',
        datasetPath: '',
        predictDatasetPath: '',
        showResults: false,
        datasetName: '',
        yamlFilePath: '',  
  
        // 控制上传类型
        uploadType: 'train',
        predictUploadType: 'predict',
  
        config: {
          epochs: 10,
          batchSize: 16
        },
        trainingStatus: 'idle', // idle, running, completed, error
        predictStatus: 'idle',
        predictProgress: 0,
        trainingResult: null,
        defaultModel: 'Default Model', // 系统默认模型
    userModels: [],               // 用户训练的模型列表，从接口获取
    availableModels: [],          // 最终下拉框使用的模型列表
    selectedModel: '',    
      };
    },
    mounted() {
    this.username = sessionStorage.getItem('username') || '';
    if (this.username) {
    this.fetchLatestTrainingResult()
  }
  },
  created() {
  this.loadUserModelsFromCache()
},
  computed: {
    // 计算属性，用于生成所有图片的 Base64 URL 列表
    allBase64Images() {
      return [
        this.trainingResult.results_png,
        this.trainingResult.p_curve,
        this.trainingResult.r_curve
      ].map(img => 'data:image/png;base64,' + img).filter(Boolean);
    },
    // 将训练结果转换为图片数组
    images() {
      return [
        { image_base64: this.trainingResult.results_png },
        { image_base64: this.trainingResult.p_curve },
        { image_base64: this.trainingResult.r_curve }
      ].filter(img => img.image_base64); // 过滤掉没有值的图片
    }
  },
    methods: {
      handleTrainingComplete(bestModelName) {
  if (!bestModelName) {
    console.warn('没有模型名可供缓存')
    return
  }

  // 防止重复添加
  if (!this.userModels.includes(bestModelName)) {
    this.userModels.unshift(bestModelName) // 插入到最前面
    localStorage.setItem('userTrainedModels', JSON.stringify(this.userModels))
    this.availableModels = [this.defaultModel, ...this.userModels]
    this.selectedModel = bestModelName // 自动选中新模型
  }
},
      loadUserModelsFromCache() {
    const cached = localStorage.getItem('userTrainedModels')
    if (cached) {
      this.userModels = JSON.parse(cached)
    } else {
      this.userModels = []
    }

    this.availableModels = [this.defaultModel, ...this.userModels]
    if (this.availableModels.length > 0) {
      this.selectedModel = this.defaultModel
    }},
      getTitle(index) {
      const titles = ['Loss & mAP 曲线', 'Precision 曲线', 'Recall 曲线'];
      return titles[index] || '';
    },
      async fetchLatestTrainingResult() {
    try {
      const res = await axios.get('http://localhost:5000/api/dcgz/train/result', {
        params: { username: this.username }
      });

      if (res.data.status === 'success') {
        this.trainingResult = res.data.result;
        this.trainingStatus = 'completed';
        this.showResults = true;
        const bestModel = this.trainingResult.best_model
        this.handleTrainingComplete(bestModel);
        this.$message.success("已恢复上次训练结果");
      }
    } catch (err) {
      console.error("获取历史训练结果失败", err);
    }
  },
      beforeUpload(file) {
        console.log(this.username)
        const isValid = file.type === 'application/zip' || file.name.endsWith('.tar.gz');
        if (!isValid) {
          this.$message.error('只能上传压缩包格式的数据集!');
          return false;
        }
        return isValid;
      },
  
      handleSuccess(response, file) {
        console.log(response)
        if(response.yaml_file) {
    this.yamlFilePath = response.yaml_file;
    this.$message.success(`上传成功，YAML 文件路径为: ${this.yamlFilePath}`
);} 
else
 {  this.$message.warning('未找到 YAML 配置文件路径'
);}
this.datasetName = response.dataset_name || file.name;
      },
  
      handlePredictSuccess(response, file) {
   
        this.predictDatasetPath = response.extracted_path || `已上传: ${file.name}`;
        console.log(this.predictDatasetPath)
        this.$message.success(`预测数据集 "${file.name}" 上传成功`);
      },
  
      async startTraining() {
  if (!this.yamlFilePath) {  // 改为使用 yamlFilePath
    this.$message.warning("请先上传数据集并生成 YAML 文件！");
    return;
  }

  this.trainingStatus = 'running';




  try {
    const res = await axios.post('http://localhost:5000/api/dcgz/train', {
      yaml_path: this.yamlFilePath,  // 新增参数
      epochs: this.config.epochs,
      batch_size: this.config.batchSize,
      username: this.username
    });
    this.pollForResult(); // 开始轮询训练结果
  } catch (err) {
    console.error(err);
    this.trainingStatus = 'error';
    this.$message.error("请求失败，请确保后端服务正常运行");
  }
},
async pollForResult(maxRetries = 20, retryInterval = 2000) {
  let retries = 0;

  const timer = setInterval(async () => {
    try {
      const res = await axios.get('http://localhost:5000/api/dcgz/train/result', {
        params: { username: this.username }
      });

      if (res.data.status === 'success') {
        clearInterval(timer);
        this.trainingResult = res.data.result;
        this.trainingStatus = 'completed';
        this.accuracy = 93.5;
        this.showResults = true;
        const bestModel = this.trainingResult.best_model
  this.handleTrainingComplete(bestModel)

        this.$message.success("训练完成");
        return;
      }

      if (++retries >= maxRetries) {
        clearInterval(timer);
        this.$message.warning("未及时获取到训练结果，请手动刷新查看");
      }
    } catch (e) {
      console.error(e);
      if (++retries >= maxRetries) clearInterval(timer);
    }
  }, retryInterval);
},
  
async startPredict() {
  if (!this.predictDatasetPath) {
    this.$message.warning("请先上传预测数据集！");
    return;
  }

  if (!this.selectedModel) {
    this.$message.warning("请选择模型文件！");
    return;
  }

  console.log("Dataset Path:", this.predictDatasetPath);
  console.log("Model Path:", this.selectedModel);

  this.predictStatus = 'running';
  try {
    const res = await axios.post('http://localhost:5000/api/dcgz/predict', {
      dataset_path: this.predictDatasetPath,
      username: this.username,
      model_path: this.selectedModel // ✅ 新增参数
    });

    if (res.data.status === 'success') {
      this.predictStatus = 'completed';
      this.showResults = true;
      this.$message.success("预测完成");
    } else {
      this.predictStatus = 'error';
      this.$message.error("预测失败：" + res.data.message);
    }
  } catch (err) {
    console.error(err);
    this.predictStatus = 'error';
    this.$message.error("请求失败，请确保后端服务正常运行");
  }
}
    }
  };
  </script>
  
  <style scoped>

  .el-card .el-card__header {
    background-color: #f5f7fa;
  }
  .el-card h4 {
    margin-top: 0;
  }
  .image-preview {
  width: 100%;
  height: auto;
  cursor: pointer;
}
.image-container {
  display: flex;
  justify-content: space-between;
}
.image-item {
  flex: 1;
  margin: 0 10px;
}
  </style>